import os
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import numpy as np
import time

# 设置环境变量以避免并行警告
os.environ['OMP_NUM_THREADS'] = '1'
os.environ['MKL_NUM_THREADS'] = '1'

# 设置PyTorch线程数
torch.set_num_threads(1)

def check_device():
    """检查并返回可用的设备（强制使用CPU）"""
    print("=== 设备检查 ===")
    device = torch.device("cpu")
    print("使用CPU设备（为确保稳定性，已禁用MPS）")
    print(f"设备: {device}")
    return device

def load_data():
    """加载并预处理MNIST数据"""
    print("\n=== 加载MNIST数据 ===")
    
    # 数据预处理
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))  # MNIST数据集的均值和标准差
    ])
    
    # 下载并加载训练数据
    trainset = torchvision.datasets.MNIST(
        root='./mnist_data', 
        train=True, 
        download=True, 
        transform=transform
    )
    trainloader = torch.utils.data.DataLoader(
        trainset, 
        batch_size=16,   # 小批处理大小
        shuffle=True,
        num_workers=0,   # 设置为0避免多进程问题
        pin_memory=False,  # 关闭pin_memory以避免问题
        persistent_workers=False  # 关闭持久化workers
    )
    
    # 下载并加载测试数据
    testset = torchvision.datasets.MNIST(
        root='./mnist_data', 
        train=False, 
        download=True, 
        transform=transform
    )
    testloader = torch.utils.data.DataLoader(
        testset, 
        batch_size=16,   # 小批处理大小
        shuffle=False,
        num_workers=0,   # 设置为0避免多进程问题
        pin_memory=False,  # 关闭pin_memory以避免问题
        persistent_workers=False  # 关闭持久化workers
    )
    
    print(f"训练集大小: {len(trainset)}")
    print(f"测试集大小: {len(testset)}")
    return trainloader, testloader

class SimpleNN(nn.Module):
    """简化版神经网络模型"""
    def __init__(self):
        super(SimpleNN, self).__init__()
        # 简单的全连接网络
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(28*28, 64)
        self.fc2 = nn.Linear(64, 32)
        self.fc3 = nn.Linear(32, 10)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.0)  # 移除dropout以提高稳定性
        
    def forward(self, x):
        x = self.flatten(x)
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.fc3(x)
        return x

def train_model(model, trainloader, testloader, device, epochs=1):
    """训练模型"""
    print(f"\n=== 开始训练 (设备: {device}) ===")
    
    # 损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.01)
    
    # 将模型移到指定设备
    model.to(device)
    print(f"模型成功迁移到 {device} 设备")
    
    # 训练过程
    for epoch in range(epochs):
        model.train()
        running_loss = 0.0
        correct = 0
        total = 0
        
        start_time = time.time()
        for i, (inputs, labels) in enumerate(trainloader):
            try:
                # 数据处理
                inputs = inputs.to(device, non_blocking=False)
                labels = labels.to(device, non_blocking=False)
                
                # 梯度清零
                optimizer.zero_grad(set_to_none=True)
                
                # 前向传播
                outputs = model(inputs)
                loss = criterion(outputs, labels)
                
                # 反向传播
                loss.backward()
                optimizer.step()
                
                # 统计信息
                running_loss += loss.item()
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
                
                # 每500个批次打印一次
                if i % 500 == 499:
                    print(f'Epoch [{epoch+1}/{epochs}], Step [{i+1}/{len(trainloader)}], '
                          f'Loss: {running_loss/500:.4f}')
                    running_loss = 0.0
                    
            except Exception as e:
                print(f"训练过程中出现错误（批次 {i}）: {e}")
                print("跳过当前批次...")
                continue
        
        epoch_time = time.time() - start_time
        
        # 计算训练准确率
        train_acc = 100 * correct / total if total > 0 else 0
        print(f'Epoch [{epoch+1}/{epochs}] 完成 - '
              f'耗时: {epoch_time:.2f}s, '
              f'训练准确率: {train_acc:.2f}%')
        
        # 评估模型
        test_acc = evaluate_model(model, testloader, device)
        print(f'测试准确率: {test_acc:.2f}%')
    
    return model

def evaluate_model(model, dataloader, device):
    """评估模型"""
    model.eval()
    correct = 0
    total = 0
    
    with torch.no_grad():
        for inputs, labels in dataloader:
            try:
                inputs = inputs.to(device, non_blocking=False)
                labels = labels.to(device, non_blocking=False)
                
                outputs = model(inputs)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
                        
            except Exception as e:
                print(f"评估时出现错误: {e}")
                continue
    
    accuracy = 100 * correct / total if total > 0 else 0
    return accuracy

def predict_samples(model, testloader, device, num_samples=3):
    """预测示例"""
    print("\n=== 预测示例 ===")
    
    model.eval()
    dataiter = iter(testloader)
    images, labels = next(dataiter)
    
    try:
        images = images.to(device, non_blocking=False)
        labels = labels.to(device, non_blocking=False)
        
        # 预测
        with torch.no_grad():
            outputs = model(images[:num_samples])
            _, predicted = torch.max(outputs, 1)
        
        # 打印预测结果
        print("前3个测试样本的预测结果:")
        for i in range(num_samples):
            print(f"样本 {i+1}: 真实标签={labels[i].item()}, 预测标签={predicted[i].item()}")
            
    except Exception as e:
        print(f"预测示例时出错: {e}")

def main():
    """主函数"""
    print("MNIST手写数字识别 - PyTorch简化神经网络实现 (CPU稳定版)")
    print("=" * 60)
    
    # 设置环境变量和线程数
    os.environ['OMP_NUM_THREADS'] = '1'
    os.environ['MKL_NUM_THREADS'] = '1'
    torch.set_num_threads(1)
    
    # 检查设备（强制使用CPU）
    device = check_device()
    
    # 加载数据
    trainloader, testloader = load_data()
    
    # 创建模型
    print("\n=== 创建简化神经网络模型 ===")
    model = SimpleNN()
    print(f"模型参数数量: {sum(p.numel() for p in model.parameters())}")
    
    # 训练模型
    try:
        model = train_model(model, trainloader, testloader, device, epochs=5)
    except Exception as e:
        print(f"训练过程中出现严重错误: {e}")
        print("程序将继续执行后续步骤...")
    
    # 预测示例
    try:
        predict_samples(model, testloader, device)
    except Exception as e:
        print(f"预测示例时出错: {e}")
    
    # 保存模型
    try:
        torch.save(model.state_dict(), '../pytorch_cpu_stable_mnist_model.pth')
        print("\n模型已保存为 'pytorch_cpu_stable_mnist_model.pth'")
    except Exception as e:
        print(f"保存模型时出错: {e}")
        print("模型未保存")
    
    print("\n" + "=" * 60)
    print("程序执行完成!")

if __name__ == "__main__":
    main()